# -*- coding: UTF-8 -*-
# Author      : WangYi
# Date        : 2020/03/16
# Description : RAW图像预处理函数 求平均 扣本底

import struct
import numpy as np
import glob
import matplotlib.pyplot as plt

'''
RAW文件读取函数
形参：raw文件 宽度和高度 其中宽度和高度用于reshape进行图像行列
返回值：np.array类型的一维数组
'''


def raw_file_read(filename, width, height):
    # read_len = width * height  # 求读出的像元
    # raw_list = list()

    # 传统读文件并转码
    # with open(filename, 'rb') as f:
    #     for i in range(read_len):
    #         tmp = f.read(2)
    #         raw_list.append(struct.unpack('H', tmp)[0])
    raw_data = np.fromfile(filename, dtype=np.uint16)

    # raw_data = raw_list # 列表数据
    # raw_data = np.array(raw_list)  # numpy数组类型的数据
    # raw_data = raw_data.reshape((height, width))  # 矩阵数据

    return raw_data


'''
# 去除帧转移效应函数
形参： file in out 输入输出文件 
width height 宽和高
integration_time 积分时间 单位us
'''
# if __name__ == "__main__":


def disposal_smearing(file_in, file_out, width, height, integration_time):
    raw_data_list = raw_file_read(file_in, width, height)
    raw_data = np.array(raw_data_list).reshape(width, height)

    t = 1.3 / integration_time

    # 生成帧转移校正矩阵
    M = np.identity(1024)  # 生成单位矩阵
    for i in range(1024):  # i 表示行
        for j in range(1024):  # j 表示列
            if j < i:
                M[i][j] = t

    M1 = np.linalg.inv(M)  # 求逆矩阵
    raw_dm = np.dot(M1, raw_data)

    raw_dm = np.clip(raw_dm, 0, 4095)  # 除去小于0的数据

    with open(file_out, 'wb') as f:
        for i in raw_dm.flat:
            foo = struct.pack('H', int(i))
            f.write(foo)

    # imgplot = plt.imshow(raw_dm, cmap='gray')
    # plt.show()
    # plt.close()


if __name__ == "__main__":
    # def cal(cnt_str):
    # 获取所有文件名
    # out_file = 'CH' + cnt_str + '-200MS.raw'
    # file_names = glob.glob('e:/1M30/CH' + cnt_str + '-200MS/*.raw')
    raw_width = 1024
    raw_height = 1024
    
    
    out_file = 'abc.raw'  # 输出文件
    file_names = glob.glob('*.raw')  # 找到该目录下所有raw文件

    # 读入所有文件数据
    raw_data = np.empty([len(file_names), raw_width*raw_height], dtype=np.uint16)
    for i, fname in enumerate(file_names):
        raw_data[i] = raw_file_read(fname, raw_width, raw_height)
    
    # all_data = list()
    # for filename in file_names:
        # all_data.append(raw_file_read(filename, 1024, 1024))

    # 形成np的数组 all_data为list类型 转为np.array型
    # raw_data = np.array(all_data)

    # 求平均
    mean_value = np.mean(raw_data, axis=0)
    # 扣本底     background subtraction
    mean_value = mean_value - 50

    # 去除负数
    res = np.clip(mean_value, 0, 4095)

    # 扣完本底的文件输出
    with open(out_file, 'wb') as f:
        for i in res:
            foo = struct.pack('H', int(i))
            f.write(foo)

    # disposal_smearing(out_file, 'kou.raw', 1024, 1024, 500)

    pass

    # imgplot = plt.imshow(dalsa_data, cmap='gray')
    # plt.show()
    # plt.close()

# if __name__ == "__main__":
#     for i in range(1, 16):
#         if i == 8:
#             pass
#         elif i == 14:
#             pass
#         else:
#             cal(str(i))
